from datetime import datetime, timedelta
import redis
from math import ceil


class Locker:

    def __init__(self, cli: redis.Redis, lock_name, lock_value, time_out=2):
        self.lock_name, self.lock_value, self.list_key = lock_name, lock_value, f'{lock_name}_list'
        self.cli, self.time_out = cli, ceil(time_out)
        self.is_locker = False

    # 加锁
    def lock(self):
        if self._lock_key(): return self.is_locker
        if self._lock_list(): return self.is_locker
        self._lock_key()
        return self.is_locker

    # 通过key获取一级锁
    def _lock_key(self):
        if self.cli.set(self.lock_name, self.lock_value, ex=self.time_out, nx=True):
            self.is_locker = True
        return self.is_locker

    # 通过队列获取二级锁
    def _lock_list(self):
        self._set_expire()
        list_top = self.cli.blpop(self.list_key, self.time_out)
        if list_top is not None:
            self.is_locker = True
            self._set_new_value()
        return self.is_locker

    # 重新设置值和过期时间
    def _set_new_value(self):
        self.cli.set(self.lock_name, self.lock_value)
        self._set_expire()

    # 重新修改过期
    def _set_expire(self):
        at = datetime.now() + timedelta(seconds=self.time_out)
        self.cli.expireat(self.lock_name, at)
        self.cli.expireat(self.list_key, at)

    # 释放锁
    def release_lock(self):
        self._set_expire()
        if self.cli.get(self.lock_name) == self.lock_value:
            self.cli.lpush(self.list_key, self.lock_value)
            self._set_expire()
